#include "flash_manager.h"
#include "base_log.h"
#include <sys/ioctl.h>
#include <mtd/mtd-user.h>
#include <fcntl.h>
#include <unistd.h>
#include <algorithm>
#include <string.h>

namespace El {
namespace Upgrade {

FlashManager::FlashManager() : fd_(-1), isOpen_(false)
{
    memset(&info_, 0, sizeof(FlashInfo));
}

FlashManager::~FlashManager()
{
    Close();
}

bool FlashManager::Open(const std::string &partition)
{
    if (isOpen_) {
        Close();
    }

    fd_ = open(partition.c_str(), O_RDWR);
    if (fd_ < 0) {
        LOG_ERROR("Failed to open MTD device: {}", partition);
        return false;
    }

    struct mtd_info_user mtd;
    if (ioctl(fd_, MEMGETINFO, &mtd) < 0) {
        LOG_ERROR("Failed to get MTD info for: {}", partition);
        Close();
        return false;
    }

    info_.size = mtd.size;
    info_.eraseSize = mtd.erasesize;
    info_.writeSize = mtd.writesize;
    info_.oobSize = mtd.oobsize;

    isOpen_ = true;
    LOG_DEBUG("Opened flash device: {}, size={}, erasesize={}, writesize={}, oobsize={}", partition, info_.size,
              info_.eraseSize, info_.writeSize, info_.oobSize);
    return true;
}

void FlashManager::Close()
{
    if (fd_ >= 0) {
        close(fd_);
        fd_ = -1;
    }
    isOpen_ = false;
}

FlashInfo FlashManager::GetInfo() const
{
    return info_;
}

bool FlashManager::Write(const uint8_t *data, size_t length, size_t offset, std::function<void(int)> progressCallback)
{
    if (!isOpen_ || !data) {
        return false;
    }

    // 检查数据是否已经对齐
    bool isAligned = (length % info_.eraseSize == 0);
    const uint8_t *writeData = data;
    size_t writeLength = length;
    std::vector<uint8_t> alignedBuffer;

    // 只有在非对齐时才创建缓冲区
    if (!isAligned) {
        // 计算需要的总大小（包含对齐填充）
        size_t alignedSize = ((length + info_.eraseSize - 1) / info_.eraseSize) * info_.eraseSize;
        // 创建对齐的缓冲区，并将数据复制到开头，后面自动填充0xFF
        alignedBuffer.resize(alignedSize, 0xFF);
        memcpy(alignedBuffer.data(), data, length);
        writeData = alignedBuffer.data();
        writeLength = alignedSize;
    }

    size_t totalWritten = 0;
    size_t currentBlock = offset / info_.eraseSize;
    size_t offsetInBlock = offset % info_.eraseSize;

    while (totalWritten < writeLength) {
        // 如果当前块已写满，移到下一个块
        if (offsetInBlock >= info_.eraseSize) {
            currentBlock++;
            offsetInBlock = 0;
        }

        // 找到好块并写入
        bool writeSuccess = false;
        while (!writeSuccess) {
            if (currentBlock >= info_.size / info_.eraseSize) {
                LOG_ERROR("Not enough good blocks, needed size: {}", length);
                return false;
            }

            loff_t blockOffset = currentBlock * info_.eraseSize;

            // 检查坏块
            if (IsBadBlock(blockOffset)) {
                LOG_WARN("Skip bad block at offset 0x{:x}", blockOffset);
                currentBlock++;
                continue;
            }

            // 如果是块的开始，需要擦除
            if (offsetInBlock == 0) {
                if (!EraseBlock(blockOffset)) {
                    LOG_WARN("Erase failed at offset 0x{:x}, marking as bad block", blockOffset);
                    currentBlock++;
                    continue;
                }
            }

            // 定位到写入位置
            loff_t writeOffset = blockOffset + offsetInBlock;
            if (lseek(fd_, writeOffset, SEEK_SET) < 0) {
                LOG_WARN("Seek failed at offset 0x{:x}", writeOffset);
                currentBlock++;
                offsetInBlock = 0;
                continue;
            }

            // 计算本次写入大小
            size_t maxWriteSize = info_.eraseSize - offsetInBlock;
            size_t remainingLength = writeLength - totalWritten;
            size_t writeSize = std::min(maxWriteSize, remainingLength);

            // 批量写入数据
            ssize_t written = write(fd_, writeData + totalWritten, writeSize);
            if (written < 0 || static_cast<size_t>(written) != writeSize) {
                LOG_WARN("Write failed at offset 0x{:x}", writeOffset + totalWritten);
                currentBlock++;
                offsetInBlock = 0;
                continue;
            }

            totalWritten += writeSize;
            offsetInBlock += writeSize;
            writeSuccess = true;

            if (progressCallback) {
                progressCallback((totalWritten * 100) / writeLength);
            }
        }
    }

    return true;
}

ssize_t FlashManager::Read(uint8_t *data, size_t length, size_t offset)
{
    if (!isOpen_ || !data) {
        return -1;
    }

    size_t totalRead = 0;
    size_t currentBlock = offset / info_.eraseSize;
    size_t offsetInBlock = offset % info_.eraseSize;

    while (totalRead < length) {
        if (currentBlock >= (info_.size / info_.eraseSize)) {
            LOG_ERROR("Read out of range");
            break;
        }

        if (IsBadBlock(currentBlock * info_.eraseSize)) {
            LOG_WARN("Skipping bad block at offset 0x{:x}", currentBlock * info_.eraseSize);
            currentBlock++;
            offsetInBlock = 0;
            continue;
        }

        loff_t readOffset = currentBlock * info_.eraseSize + offsetInBlock;
        if (lseek(fd_, readOffset, SEEK_SET) < 0) {
            LOG_ERROR("Seek failed at offset 0x{:x}", readOffset);
            break;
        }

        size_t chunkSize = std::min(info_.eraseSize - offsetInBlock, length - totalRead);
        ssize_t bytesRead = read(fd_, data + totalRead, chunkSize);
        if (bytesRead <= 0) {
            break;
        }
        totalRead += bytesRead;
        offsetInBlock += bytesRead;
        if (offsetInBlock >= info_.eraseSize) {
            currentBlock++;
            offsetInBlock = 0;
        }
    }

    return totalRead;
}

bool FlashManager::IsBadBlock(size_t blockOffset)
{
    loff_t offset = blockOffset;
    return ioctl(fd_, MEMGETBADBLOCK, &offset) > 0;
}

bool FlashManager::EraseBlock(size_t blockOffset)
{
    struct erase_info_user erase;
    erase.start = blockOffset;
    erase.length = info_.eraseSize;
    return ioctl(fd_, MEMERASE, &erase) == 0;
}

} // namespace Upgrade
} // namespace El